import os
import sys
import numpy as np
from PIL import Image

def yuv_to_rgb(yuv_path, rgb_path=None):
    """
    将YUV图像转换为RGB图像。
    参数：
        yuv_path: str，YUV图像的路径。
        rgb_path: str，可选，RGB图像的输出路径。如果未提供，则使用与YUV图像相同的文件名和扩展名。
        返回值：无
    """
    if rgb_path is None:
        rgb_path = os.path.splitext(yuv_path)[0] + '_rgb' + '.png'
    
    # 打开YUV图像并转换为numpy数组
    yuv_img = Image.open(yuv_path)
    yuv_array = np.array(yuv_img)
    
    # 处理通道格式
    if len(yuv_array.shape) == 2:
        yuv_array = np.stack([yuv_array]*3, axis=-1)
    elif yuv_array.shape[2] > 3:
        yuv_array = yuv_array[:, :, :3]
    
    # 转换为float32类型并归一化
    yuv_float = yuv_array.astype(np.float32) / 255.0
    
    # 分离YUV通道（ITU-R BT.601标准）
    Y = yuv_float[..., 0]          # Y分量
    U = yuv_float[..., 1] - 0.5    # U分量偏移
    V = yuv_float[..., 2] - 0.5    # V分量偏移
    
    # 应用ITU-R BT.601转换公式
    R = Y + 1.403 * V
    G = Y - 0.344 * U - 0.714 * V
    B = Y + 1.773 * U
    
    # 合并通道并裁剪
    rgb_float = np.stack([R, G, B], axis=-1)
    rgb_float = np.clip(rgb_float, 0, 1)
    
    # 转换回uint8格式
    rgb_array = (rgb_float * 255).astype(np.uint8)
    
    # 保存RGB图像
    rgb_img = Image.fromarray(rgb_array)
    print("Saving RGB image to", rgb_path)
    rgb_img.save(rgb_path)

if __name__ == '__main__':
    if len(sys.argv) > 2:
        yuv_to_rgb(sys.argv[1], sys.argv[2])
    elif len(sys.argv) == 2:
        yuv_to_rgb(sys.argv[1])
    else:
        raise ValueError('Invalid number of arguments. Usage: python yuv_to_rgb.py <yuv_file> [output_file]')
